import base64
import streamlit as st
from PIL import Image
import io
from large_model.agent_executors import Executors
from streamlit_model.streamlit_utils.session_utils import load_session, save_session

PAGE_KEY = "uploaded_img_case"

def uploaded_img_case():
    cached = load_session(PAGE_KEY)

    # 状态初始化（去掉 img_data 和 img_b64 的缓存）
    for k in ["img_data", "img_b64", "exec_result"]:
        if k not in st.session_state:
            st.session_state[k] = None  # 不再从缓存恢复图片

    _, title_col, _ = st.columns([1, 3, 1])
    with title_col:
        st.title("上传图片生成用例")

    _, uploader_col, _ = st.columns([1, 3, 1])
    with uploader_col:
        uploaded_file = st.file_uploader(
            "请上传图片",
            type=["png", "jpg", "jpeg", "bmp", "gif"],
            label_visibility="collapsed"
        )

        # 上传图片时只保留在 session_state，不写入缓存
        if uploaded_file is not None:
            st.session_state.img_data = uploaded_file.read()
            st.session_state.img_b64 = base64.b64encode(st.session_state.img_data).decode()
            st.session_state.exec_result = None
            st.success("✅ 图片上传成功！")

    _, desc_col, _ = st.columns([1, 3, 1])
    with desc_col:
        description = st.text_area(
            "请输入描述",
            height=300,
            value=cached.get("description", "需要编写测试用例的功能或者需求:"),
            max_chars=3000
        )

    _, btn_col, _ = st.columns([2, 1, 2])
    with btn_col:
        if st.button("▶️ 生成测试用例", use_container_width=True):
            if st.session_state.img_b64 is None:
                st.warning("请先上传图片！")
            else:
                with st.spinner("正在执行测试用例..."):
                    result = Executors().visual_main(description, st.session_state.img_b64)
                    st.session_state.exec_result = result

    # 即时保存
    # 保存缓存时，不再保存 img_data 和 img_b64
    save_session(PAGE_KEY, {
        "description": description,
        "exec_result": st.session_state.exec_result
        # 不保存 img_data 和 img_b64
    })

    if st.session_state.exec_result:
        _, res_col, _ = st.columns([1, 3, 1])
        with res_col:
            st.subheader("执行结果")
            st.success(st.session_state.exec_result)